from typing import Union
from fastapi import FastAPI
import uvicorn
import psutil
import os
import argparse

app = FastAPI()


def kill_task(process_name:str):
    pids = psutil.pids()
    for pid in pids:
        try:
            p = psutil.Process(pid)
        except Exception as e:
            continue
        # print('pid-%s,pname-%s' % (pid, p.name()))
        if process_name in p.name():
            cmd = 'taskkill /F /IM {}'.format(p.name())
            print(cmd)
            os.system(cmd)

@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.get("/kill_process/{process_name}")
def kill_process(process_name: str):
    kill_task(process_name)
    return {"status":"ok"}

@app.get("/reboot")
def reboot():
    os.system("shutdown /r /t 1")
    return {"status":"ok"}

@app.get("/shutdown")
def shutdown():
    os.system("shutdown /s /t 1")
    return {"status":"ok"}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MSCTXRC远程控制程序")
    parser.add_argument('-m',"--mode", type = str, help="运行模式，默认为，可选web为cmd",default="web")
    parser.add_argument('-t',"--task", type = str, help="任务,k为杀死进程名，r为重启，s为关机",required=False)    
    parser.add_argument('-kpn',"--kill_process_name", type = str, help="杀死进程名",required=False)
    args = parser.parse_args()
    mode = args.mode
    if mode == "web":
        uvicorn.run("main:app", host="0.0.0.0", port=18636, reload=False)
    else:
        if args.task == "k":
            kill_task(args.kill_process_name)
        elif args.task == "r":
            reboot()
        elif args.task == "s":
            shutdown()
    
